import streamlit as st
import os
import json
import pandas as pd
from PIL import Image
import time
from datetime import datetime
import matplotlib.pyplot as plt
import seaborn as sns
from io import BytesIO
import base64

from models.vllm_model import VLLMModel
from utils.helpers import save_upload_file

def show_detection():
    """显示缺陷检测页面"""
    
    st.title("缺陷检测")
    
    # 初始化会话状态
    if "vllm_model" not in st.session_state:
        st.session_state.vllm_model = None
    
    # 侧边栏配置
    st.sidebar.header("模型配置")
    
    # Ollama服务地址配置
    ollama_url = st.sidebar.text_input("Ollama服务地址", value="http://localhost:11434")
    
    # 模型选择
    try:
        if st.session_state.vllm_model is None:
            # 初始化VLLM模型
            st.session_state.vllm_model = VLLMModel(base_url=ollama_url)
        
        # 刷新模型列表
        if st.sidebar.button("刷新模型列表"):
            refresh_placeholder = st.sidebar.empty()
            refresh_placeholder.info("正在获取可用模型...")
            try:
                available_models = st.session_state.vllm_model.get_available_models()
                st.session_state.available_models = available_models
                refresh_placeholder.success("模型列表已更新")
            except Exception as e:
                refresh_placeholder.error("获取模型列表失败")
    except Exception as e:
        st.sidebar.error(f"连接Ollama服务失败: {str(e)}")
        st.sidebar.warning("请确保Ollama服务已启动，并检查服务地址是否正确")
        available_models = ["qwen2.5-vl"]
    
    # 如果会话状态中有模型列表，则使用它
    if "available_models" in st.session_state:
        available_models = st.session_state.available_models
    else:
        available_models = ["qwen2.5-vl"]
    
    # 选择模型
    selected_model = st.sidebar.selectbox("选择VLLM模型", available_models)
    
    # 更新模型配置
    if st.sidebar.button("应用配置"):
        try:
            st.session_state.vllm_model = VLLMModel(base_url=ollama_url, model_name=selected_model)
            st.sidebar.success(f"成功配置模型: {selected_model}")
        except Exception as e:
            st.sidebar.error(f"配置模型失败: {str(e)}")
    
    # 主界面
    st.header("上传图片进行缺陷检测")
    
    # 检查知识库和CLIP模型
    if "knowledge_base" not in st.session_state or st.session_state.knowledge_base is None:
        st.warning("请先在'知识库管理'页面加载CLIP模型和知识库")
        return
    
    # 获取可用的知识库
    kb_list = st.session_state.knowledge_base.get_product_list()
    
    if not kb_list:
        st.warning("未找到任何知识库，请先创建知识库")
        return
    
    # 选择产品
    selected_product = st.selectbox("选择产品", kb_list)
    
    # 上传图片
    uploaded_file = st.file_uploader("上传要检测的图片", type=["png", "jpg", "jpeg"])
    
    # 显示检测结果
    if uploaded_file is not None:
        # 显示上传的图片
        image = Image.open(uploaded_file)
        st.image(image, caption="上传的图片", use_column_width=True)
        
        # 保存上传的图片
        upload_dir = os.path.join("./uploads", "detection")
        os.makedirs(upload_dir, exist_ok=True)
        img_path = save_upload_file(uploaded_file, upload_dir)
        
        # 加载知识库
        if st.button("进行检测"):
            with st.spinner("正在检测..."):
                try:
                    # 加载知识库
                    st.session_state.knowledge_base.load_or_create(selected_product)
                    
                    # 从知识库检索相似图片
                    retrieval_results = st.session_state.knowledge_base.search(img_path, top_k=5)
                    
                    # 调用VLLM模型进行检测
                    result = st.session_state.vllm_model.detect_defect(
                        image_path=img_path,
                        product_name=selected_product,
                        knowledge_results=retrieval_results
                    )
                    
                    # 获取并保存prompt
                    retrieval_context = st.session_state.vllm_model._format_retrieval_results(retrieval_results)
                    prompt = f"""你是一个产品表面缺陷检测专家，请仔细分析图片中的{selected_product}产品，并检测是否存在表面缺陷。
        
以下是从相似图片知识库中检索到的结果供你参考:
{retrieval_context}

请分析图片并输出标准JSON格式结果，包括:
1. has_defect: 布尔值，表示是否存在缺陷
2. defect_type: 如存在缺陷，提供缺陷类型；如无缺陷，则为"good"或"无"
3. confidence: 置信度，0-1的浮点数
4. analysis: 简短分析说明（不超过100字）

只需要输出JSON格式结果，不要有其他回复。"""
                    
                    # 显示检索结果
                    st.header("知识库检索结果")
                    
                    # 创建结果表格
                    retrieval_df = []
                    
                    # 添加检索结果
                    for i, res in enumerate(retrieval_results):
                        retrieval_df.append({
                            "序号": i + 1,
                            "图片路径": res.get("path", ""),
                            "标签": res.get("label", "未知"),
                            "相似度": f"{res.get('similarity', 0):.4f}"
                        })
                    
                    # 显示表格
                    if retrieval_df:
                        df = pd.DataFrame(retrieval_df)
                        st.dataframe(df, use_container_width=True)
                        
                        # 显示检索到的图片
                        st.subheader("检索到的相似图片")
                        cols = st.columns(min(5, len(retrieval_results)))
                        for i, res in enumerate(retrieval_results[:5]):  # 最多显示5张图片
                            try:
                                img_path = res.get("path", "")
                                if os.path.exists(img_path):
                                    with cols[i]:
                                        img = Image.open(img_path)
                                        st.image(img, caption=f"{res.get('label', '未知')} ({res.get('similarity', 0):.4f})", use_column_width=True)
                            except Exception as e:
                                st.error(f"无法加载图片 {img_path}: {str(e)}")
                    else:
                        st.warning("未找到相似图片")
                    
                    # 显示传给大模型的prompt
                    with st.expander("查看传给大模型的Prompt"):
                        st.code(prompt, language="markdown")
                    
                    # 显示检测结果
                    st.header("检测结果")
                    
                    if result.get("success", False):
                        response = result.get("response", {})
                        
                        # 如果是字符串，尝试解析JSON
                        if isinstance(response, str):
                            try:
                                # 尝试提取JSON部分
                                start = response.find("{")
                                end = response.rfind("}") + 1
                                if start >= 0 and end > start:
                                    json_str = response[start:end]
                                    response = json.loads(json_str)
                                    st.session_state.last_detection = response
                                else:
                                    st.text_area("原始回复", value=response, height=200)
                                    st.warning("无法解析为JSON格式")
                                    return
                            except:
                                st.text_area("原始回复", value=response, height=200)
                                st.warning("无法解析为JSON格式")
                                return
                        
                        # 保存到会话状态
                        st.session_state.last_detection = response
                        
                        # 提取结果
                        has_defect = response.get("has_defect", False)
                        defect_type = response.get("defect_type", "未知")
                        confidence = response.get("confidence", 0.0)
                        analysis = response.get("analysis", "")
                        
                        # 结果展示
                        col1, col2 = st.columns(2)
                        
                        # 左侧展示文本结果
                        with col1:
                            # 显示结果信息
                            if has_defect:
                                st.error(f"检测结果: 存在缺陷")
                            else:
                                st.success(f"检测结果: 无缺陷")
                            
                            st.info(f"缺陷类型: {defect_type}")
                            st.info(f"置信度: {confidence:.4f}")
                            st.text_area("分析说明", value=analysis, height=150)
                        
                        # 右侧展示可视化结果
                        with col2:
                            # 创建置信度条形图
                            fig, ax = plt.subplots(figsize=(5, 3))
                            
                            # 设置中文字体
                            plt.rcParams['font.sans-serif'] = ['SimHei']  
                            plt.rcParams['axes.unicode_minus'] = False
                            
                            if has_defect:
                                ax.barh("缺陷", confidence, color='red')
                            else:
                                ax.barh("正常", confidence, color='green')
                                
                            ax.set_xlim(0, 1)
                            ax.set_xlabel("置信度")
                            ax.grid(axis='x', linestyle='--', alpha=0.7)
                            st.pyplot(fig)
                        
                        # 显示原始JSON
                        with st.expander("查看原始JSON"):
                            st.json(response)
                        
                        # 保存结果
                        result_dir = os.path.join("./results", selected_product)
                        os.makedirs(result_dir, exist_ok=True)
                        
                        result_file = os.path.join(result_dir, f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_result.json")
                        with open(result_file, 'w', encoding='utf-8') as f:
                            json.dump({
                                "timestamp": datetime.now().isoformat(),
                                "product": selected_product,
                                "image_path": img_path,
                                "result": response,
                                "retrieval": [{
                                    "path": res.get("path", ""),
                                    "label": res.get("label", ""),
                                    "similarity": float(res.get("similarity", 0))
                                } for res in retrieval_results]
                            }, f, ensure_ascii=False, indent=2)
                        
                        st.success(f"检测结果已保存至 {result_file}")
                    else:
                        st.error(f"检测失败: {result.get('error', '未知错误')}")
                
                except Exception as e:
                    st.error(f"检测过程发生错误: {str(e)}")
    
    st.markdown("---")
    
    st.markdown("""
    ### 使用说明
    
    1. 选择已创建的产品知识库
    2. 上传要检测的图片
    3. 点击"进行检测"按钮
    4. 系统会返回检测结果，包括:
       - 是否存在缺陷
       - 缺陷类型
       - 置信度
       - 分析说明
    
    检测结果会自动保存，可在 ./results/{产品名} 目录中找到。
    """) 